import cv2
import mediapipe as mp
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, Ellipse

# 初始化 MediaPipe Pose
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5)

# 加载图像并进行姿态检测
image_path = 'your_image_path.jpg'  # 替换为您的图片路径
image = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转为 RGB 格式
results = pose.process(image_rgb)

# 创建透明背景
fig, ax = plt.subplots(figsize=(6, 10), dpi=100)
ax.set_aspect('equal')
ax.set_facecolor((0, 0, 0, 0))  # 透明背景

# 获取关键点位置
if results.pose_landmarks:
    landmarks = results.pose_landmarks.landmark

    # 定义头部和身体部分的连接
    head_index = 0  # 通常头部的索引为 0 (需要根据 MediaPipe 定义)
    connections = mp_pose.POSE_CONNECTIONS

    # 绘制头部为圆形
    head = landmarks[head_index]
    head_x, head_y = head.x * image.shape[1], head.y * image.shape[0]
    head_circle = Circle((head_x, head_y), 20, color='white')  # 半径可根据图像大小调整
    ax.add_patch(head_circle)

    # 绘制其他身体连接部分为椭圆型线段
    for connection in connections:
        start_idx, end_idx = connection
        start = landmarks[start_idx]
        end = landmarks[end_idx]
        
        # 获取连接点坐标
        start_x, start_y = start.x * image.shape[1], start.y * image.shape[0]
        end_x, end_y = end.x * image.shape[1], end.y * image.shape[0]
        
        # 计算两点之间的中点和角度
        center_x, center_y = (start_x + end_x) / 2, (start_y + end_y) / 2
        dx, dy = end_x - start_x, end_y - start_y
        angle = np.degrees(np.arctan2(dy, dx))  # 角度计算
        
        # 创建椭圆形
        length = np.sqrt(dx**2 + dy**2)  # 椭圆的长度为两点距离
        ellipse = Ellipse((center_x, center_y), length, 10, angle=angle, color='white')  # 宽度可根据需求调整
        ax.add_patch(ellipse)

# 隐藏坐标轴
plt.axis('off')

# 保存为带透明背景的图片
plt.savefig('skeleton_transparent.png', format='png', transparent=True)
plt.show()